

# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

import numpy as np
import matplotlib.pyplot as plt

from kalman_filter._01_kf_point_mass_model import PointMassModelKalmanFilter

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
plt.rc("font", size=15)

def generate_control_sequence(time_steps, pattern="acceleration_deceleration"):
    """
    生成控制序列
    
    参数:
    time_steps: 时间步数
    pattern: 控制模式
    
    返回:
    control_sequence: 控制序列
    """
    control_sequence = []
    
    if pattern == "acceleration_deceleration":
        # 加速-减速模式
        for t in range(time_steps):
            if t < time_steps // 5:
                control_sequence.append(1.0)  # 加速
            elif t < 2 * time_steps // 5:
                control_sequence.append(-0.5)  # 减速
            elif t < 3 * time_steps // 5:
                control_sequence.append(0.0)  # 匀速
            elif t < 4 * time_steps // 5:
                control_sequence.append(-2.0)  # 急刹车
            else:
                control_sequence.append(0.5)  # 再加速
                
    elif pattern == "sinusoidal":
        # 正弦波模式
        for t in range(time_steps):
            control_sequence.append(2.0 * np.sin(2 * np.pi * t / time_steps))
    
    elif pattern == "step":
        # 阶跃模式
        for t in range(time_steps):
            if t < time_steps // 3:
                control_sequence.append(1.0)
            elif t < 2 * time_steps // 3:
                control_sequence.append(-1.0)
            else:
                control_sequence.append(0.0)
    
    else:
        # 默认：零控制
        control_sequence = [0.0] * time_steps
    
    return control_sequence

def simulate_vehicle_motion(duration, dt):
    """
    模拟车辆运动生成测试数据
    """
    time_steps = int(duration / dt)
    t = np.linspace(0, duration, time_steps)

    # 真实状态
    true_states = np.zeros((time_steps, 3))
    # 初始状态
    true_states[0] = [0, 5, 0]  # [y, v, a]

    # 生成控制序列 [加速度]
    controls = generate_control_sequence(time_steps, pattern="sinusoidal")

    # 模拟车辆运动（基于控制序列迭代车辆每一帧状态）
    for i in range(1, time_steps):
        # 简单的欧拉积分
        y, v, a = true_states[i-1]
        a_control = controls[i-1]

        v_new = v + a_control * dt
        y_new = y + v_new * dt + 0.5 * a_control * dt**2

        true_states[i] = [y_new, v_new, a_control]

    return t, true_states, controls

def add_noise(data, stds_dev, noise_type="gaussian"):
    """添加噪声到数据"""
    stds_dev = np.array(stds_dev)
    if noise_type == "gaussian":
        # 高斯噪声
        noise = np.random.normal(0, stds_dev, data.shape)
    elif noise_type == "uniform":
        # 均匀噪声范围为 [-std, std]
        noise = np.random.uniform(-stds_dev, stds_dev, data.shape)
    return data + noise

def compare_state(true_states, measurements):
    for i in range(int(true_states.shape[0] / 5)):
        print("y状态: ", true_states[i, 0], measurements[i, 0])
        print("v状态: ", true_states[i, 1], measurements[i, 1])
        print("a状态: ", true_states[i, 2], measurements[i, 2])
        print("")

if __name__ == "__main__":
  # 参数设置
  dt = 0.1  # 时间步长
  duration = 10  # 模拟时长

  # 生成测试数据
  t, true_states, controls = simulate_vehicle_motion(duration, dt)

  # 添加噪声创建测量值
  measurement_noise_std = [0.5, 0.5, 0.1]
  measurements = add_noise(true_states, measurement_noise_std,  noise_type="gaussian")
#   compare_state(true_states, measurements)

  # 创建卡尔曼滤波器
  kf = PointMassModelKalmanFilter(dt=dt,
                           process_noise_std=[0.1, 0.1, 0.8],
                           measurement_noise_std=measurement_noise_std)

  # 设置初始状态
  kf.reset(np.array(true_states[0]))  # 初始状态估计
  # 存储估计结果
  estimated_states = np.zeros_like(true_states)
  estimated_states[0] = kf.get_state()

  # 存储每个状态的协方差
  estimated_covariances = np.zeros((len(t), kf.state_dim))

  # 运行卡尔曼滤波
  for i in range(1, len(t)):
      kf.filter_step(measurements[i], controls[i-1])
      estimated_states[i] = kf.get_state()
      # 保存协方差信息
      P = kf.get_covariance()
      estimated_covariances[i] = [P[0,0], P[1,1], P[2,2]]  # 只存储对角线元素（方差）

  # 绘制结果
  plt.figure(figsize=(12, 8))

  # Y位置随时间变化
  plt.subplot(2, 3, 1)
  plt.plot(t, true_states[:, 0], 'g-', label='真实值', linewidth=2)
  plt.plot(t, measurements[:, 0], 'ro', markersize=2, alpha=0.5, label='测量值')
  plt.plot(t, estimated_states[:, 0], 'b-', label='估计值', linewidth=1)
  plt.xlabel('时间 (s)')
  plt.ylabel('Y位置 (m)')
  plt.title('Y位置')
  plt.legend()
  plt.grid(True)

  # 速度随时间变化
  plt.subplot(2, 3, 2)
  plt.plot(t, true_states[:, 1], 'g-', label='真实值', linewidth=2)
  plt.plot(t, measurements[:, 1], 'ro', markersize=2, alpha=0.5, label='测量值')
  plt.plot(t, estimated_states[:, 1], 'b-', label='估计值', linewidth=1)
  plt.xlabel('时间 (s)')
  plt.ylabel('速度 (m/s)')
  plt.title('速度')
  plt.legend()
  plt.grid(True)

  # 加速度随时间变化
  plt.subplot(2, 3, 3)
  plt.plot(t, true_states[:, 2], 'g-', label='真实值', linewidth=2)
  plt.plot(t, measurements[:, 2], 'ro', markersize=2, alpha=0.5, label='测量值')
  plt.plot(t, estimated_states[:, 2], 'b-', label='估计值', linewidth=1)
  plt.xlabel('时间 (s)')
  plt.ylabel('速度 (m/s^2)')
  plt.title('加速度')
  plt.legend()
  plt.grid(True)

  # 每个状态的协方差
  plt.subplot(2, 3, 4)
  plt.plot(t, estimated_covariances[:, 0], label='位置', linewidth=1)
  plt.plot(t, estimated_covariances[:, 1], label='速度', linewidth=1)
  plt.plot(t, estimated_covariances[:, 2], label='加速度', linewidth=1)
  plt.xlabel('时间 (s)')
  plt.ylabel('方差')
  plt.title('状态方差')
  plt.legend()
  plt.grid(True)

  plt.tight_layout()
  plt.show()